#include <elf.h>
#include <stddef.h>
#include <types.h>
#include <stdio.h>
#include <malloc.h>

static bool_t is_elf_magic(unsigned char e_ident[EI_NIDENT])
{
    return (e_ident[EI_MAG0] == ELFMAG0 &&
            e_ident[EI_MAG1] == ELFMAG1 &&
            e_ident[EI_MAG2] == ELFMAG2 && e_ident[EI_MAG3] == ELFMAG3);
}

static bool_t parse_elf_header(const char *code, Elf64_Ehdr *header)
{
    if (!is_elf_magic(code))
    {
        printf("is not elf magic\n");
        return FALSE;
    }
    if (code[EI_CLASS] == ELFCLASSNONE)
    {
        printf("Invalid class\n");
        return FALSE;
    }
    if (code[EI_CLASS] == ELFCLASS32)
    {
        printf("32-bit objects\n");
        return FALSE;
    }
    if (code[EI_CLASS] == ELFCLASS64)
    {
        printf("64-bit objects\n");
    }
    if (code[EI_DATA] == ELFDATANONE)
    {
        printf("Invalid data encoding\n");
        return FALSE;
    }
    if (code[EI_DATA] == ELFDATA2LSB)
    {
        printf("2's complement, little endian\n");
    }
    if (code[EI_DATA] == ELFDATA2MSB)
    {
        printf("2's complement, big endian\n");
        return FALSE;
    }

    *header = *(Elf64_Ehdr *)code;

    if (header->e_machine != EM_RISCV)
    {
        printf("e_machine is not RISC-V\n");
        return FALSE;
    }

    return TRUE;
}

bool_t elf_parse_file(const char *code)
{
    Elf64_Ehdr ehdr;

    if (!parse_elf_header(code, &ehdr))
        return FALSE;

    printf("e_type:%hu\t\tObject file type\n", ehdr.e_type);                          /* Object file type */
    printf("e_machine:%hu\t\tArchitecture\n", ehdr.e_machine);                        /* Architecture */
    printf("e_version:%u\t\tObject file version\n", ehdr.e_version);                  /* Object file version */
    printf("e_entry:%lx\t\tEntry point virtual address\n", ehdr.e_entry);             /* Entry point virtual address */
    printf("e_phoff:%lx\t\tProgram header table file offset\n", ehdr.e_phoff);        /* Program header table file offset */
    printf("e_shoff:%lx\t\tSection header table file offset\n", ehdr.e_shoff);        /* Section header table file offset */
    printf("e_flags:%u\t\tProcessor-specific flags\n", ehdr.e_flags);                 /* Processor-specific flags */
    printf("e_ehsize:%hu\t\tELF header size in bytes\n", ehdr.e_ehsize);              /* ELF header size in bytes */
    printf("e_phentsize:%hu\t\tProgram header table entry size\n", ehdr.e_phentsize); /* Program header table entry size */
    printf("e_phnum:%hu\t\tProgram header table entry count\n", ehdr.e_phnum);        /* Program header table entry count */
    printf("e_shentsize:%hu\t\tSection header table entry size\n", ehdr.e_shentsize); /* Section header table entry size */
    printf("e_shnum:%hu\t\tSection header table entry count\n", ehdr.e_shnum);        /* Section header table entry count */
    printf("e_shstrndx:%hu\t\tSection header string table index\n", ehdr.e_shstrndx); /* Section header string table index */

    /* Parse program headers and section headers */
    for (int i = 0; i < ehdr.e_phnum; ++i)
    {
        printf("Elf64_Phdr[%d]\n", i);

        Elf64_Phdr *phdr = code + ehdr.e_phoff + ehdr.e_phentsize * i;

        printf("p_type:%u\t\tSegment type\n", phdr->p_type);                /* Segment type */
        printf("p_flags:%x\t\tSegment flags\n", phdr->p_flags);             /* Segment flags */
        printf("p_offset:%lx\t\tSegment file offset\n", phdr->p_offset);    /* Segment file offset */
        printf("p_vaddr:%lx\t\tSegment virtual address\n", phdr->p_vaddr);  /* Segment virtual address */
        printf("p_paddr:%lx\t\tSegment physical address\n", phdr->p_paddr); /* Segment physical address */
        printf("p_filesz:%lu\t\tSegment size in file\n", phdr->p_filesz);   /* Segment size in file */
        printf("p_memsz:%lu\t\tSegment size in memory\n", phdr->p_memsz);   /* Segment size in memory */
        printf("p_align:%lu\t\tSegment alignment\n", phdr->p_align);        /* Segment alignment */
    }

    for (int i = 0; i < ehdr.e_shnum; ++i)
    {
        printf("Elf64_Shdr[%d]\n", i);

        Elf64_Shdr *shdr = code + ehdr.e_shoff + ehdr.e_shentsize * i;
        printf("sh_name:%u\t\tSection name (string tbl index)\n", shdr->sh_name);          /* Section name (string tbl index) */
        printf("sh_type:%u\t\tSection type\n", shdr->sh_type);                             /* Section type */
        printf("sh_flags:%lx\t\tSection flags\n", shdr->sh_flags);                         /* Section flags */
        printf("sh_addr:%lx\t\tSection virtual addr at execution\n", shdr->sh_addr);       /* Section virtual addr at execution */
        printf("sh_offset:%lx\t\tSection file offset\n", shdr->sh_offset);                 /* Section file offset */
        printf("sh_size:%lu\t\tSection size in bytes\n", shdr->sh_size);                   /* Section size in bytes */
        printf("sh_link:%u\t\tLink to another section\n", shdr->sh_link);                  /* Link to another section */
        printf("sh_info:%u\t\tAdditional section information\n", shdr->sh_info);           /* Additional section information */
        printf("sh_addralign:%lu\t\tSection alignment\n", shdr->sh_addralign);             /* Section alignment */
        printf("sh_entsize:%lu\t\tEntry size if section holds table\n", shdr->sh_entsize); /* Entry size if section holds table */
    }
    return TRUE;
}

struct elf_info_t
{
    Elf64_Ehdr ehdr;
    Elf64_Phdr *phdr;
    Elf64_Shdr *shdr;
};

struct elf_info_t *elf_get_info(const char *code)
{
    Elf64_Ehdr ehdr;

    if (!parse_elf_header(code, &ehdr))
        return NULL;

    struct elf_info_t *info = malloc(sizeof(struct elf_info_t));
    info->ehdr = ehdr;
    info->phdr = malloc(sizeof(Elf64_Phdr) * ehdr.e_phnum);
    info->shdr = malloc(sizeof(Elf64_Shdr) * ehdr.e_shnum);

    /* Parse program headers and section headers */
    for (int i = 0; i < ehdr.e_phnum; ++i)
    {
        info->phdr[i] = *(Elf64_Phdr *)(code + ehdr.e_phoff + ehdr.e_phentsize * i);
    }

    for (int i = 0; i < ehdr.e_shnum; ++i)
    {
        info->shdr[i] = *(Elf64_Shdr *)(code + ehdr.e_shoff + ehdr.e_shentsize * i);
    }
    return info;
}

void elf_free_info(struct elf_info_t *info)
{
    if (info)
    {
        free(info->phdr);
        free(info->shdr);
        free(info);
    }
}